import copy
import math

import common_modules

from DrawingAPIs.breakdown_drawer import *
from DrawingAPIs.bar_drawer import *

field_names = ['Device', 'DNN model', 'Original', 'IntraLayer', 'Pipeline',
               'IntraLayer-Speedup', 'Pipeline-Speedup']
NN_MODELS = ['CIFAR10_nv', 'AlexNet', 'Resnet18', 'VGG16']

def PipeSpeedupFileProcess(pipe_spdup_file, devices):
    pipe_speedup_data = {}
    for net in NN_MODELS:
        pipe_speedup_data.setdefault(net, [])

    org_pipe_vals = {}
    with open(pipe_spdup_file, 'r') as pipe_reader:
        pipe_csv = csv.DictReader(pipe_reader)

        for row in pipe_csv:
            curr_dev = row['Device']
            if curr_dev not in org_pipe_vals.keys():
                org_pipe_vals.setdefault(curr_dev, {})
            curr_model = row['DNN model']
            if curr_model not in org_pipe_vals[curr_dev].keys():
                org_pipe_vals[curr_dev].setdefault(curr_model, {})
            org_pipe_vals[curr_dev][curr_model]['IntraLayer'] = round(float(row['IntraLayer-Speedup']), 3)
            org_pipe_vals[curr_dev][curr_model]['Pipeline'] = round(float(row['Pipeline-Speedup']), 3) - \
                round(float(row['IntraLayer-Speedup']), 3)

    for model in NN_MODELS:
        intra_dev_vals = []
        pipe_dev_vals = []
        for dev in devices:
            intra_dev_vals.append(org_pipe_vals[dev][model]['IntraLayer'])
            pipe_dev_vals.append(org_pipe_vals[dev][model]['Pipeline'])
        pipe_speedup_data[model] = [intra_dev_vals, pipe_dev_vals]

    return pipe_speedup_data

def OverallSpeedupFileProcess(pipe_spdup_file, devices):
    org_pipe_vals = {}
    with open(pipe_spdup_file, 'r') as pipe_reader:
        pipe_csv = csv.DictReader(pipe_reader)

        for row in pipe_csv:
            curr_dev = row['Device']
            if curr_dev not in org_pipe_vals.keys():
                org_pipe_vals.setdefault(curr_dev, {})
            curr_model = row['DNN model']
            if curr_model not in org_pipe_vals[curr_dev].keys():
                org_pipe_vals[curr_dev].setdefault(curr_model, {})
            org_pipe_vals[curr_dev][curr_model]['IntraLayer'] = round(float(row['IntraLayer-Speedup']), 3)
            org_pipe_vals[curr_dev][curr_model]['Pipeline'] = round(float(row['Pipeline-Speedup']), 3)

    overall_speedups = []
    y_max_val = 0.0
    for dev in org_pipe_vals.keys():
        dev_speedup = []
        for model in NN_MODELS:
            dev_speedup.append(org_pipe_vals[dev][model]['Pipeline'])
        y_max_val = max(y_max_val, max(dev_speedup))
        print('Dev_speedup = {0}'.format(dev_speedup))
        overall_speedups.append(dev_speedup)

    return overall_speedups, y_max_val

def DrawOverallSpeedupFig(args, devices):
    pipeline_speedup_file = args.root[0] + '/Pipeline/Pipeline_Speedup.csv'
    pipe_speedup_data = PipeSpeedupFileProcess(pipeline_speedup_file, devices)
    overall_speedups, y_max_val = OverallSpeedupFileProcess(pipeline_speedup_file, devices)

    # legend_labels = ['IntraLayer', 'HGP4NN']
    # for net in NN_MODELS:
    #     y_max_val = max(pipe_speedup_data[net][0]) + max(pipe_speedup_data[net][1])
    #     print("Maximum speedup with HGP4NN: {0}".format(y_max_val))
    #     DrawBreakdownFig(devices, pipe_speedup_data[net], legend_labels, x_label = '', y_label = 'Speedup',
    #                      x_min = 0.5, x_max = (len(devices) + 0.5), y_min = 0, y_max = y_max_val, bar_width = 0.4,
    #                      folder = args.output[0] + '/Speedup/Overall', filename = 'exp_overall_speedup_' + net)
    bar_width = 0.45
    x_ticks = np.arange(2, 2 * (len(overall_speedups[0]) + 1), 2)
    print('Before ceil operation y_max_val = {0}'.format(y_max_val))
    if y_max_val > (math.ceil(y_max_val) - 0.5):
        y_max_val = math.ceil(y_max_val)
    else:
        y_max_val = math.ceil(y_max_val) - 0.5

    print('y_max_val = {0}'.format(y_max_val))
    DrawBarFig(x_ticks, x_ticklabels=NN_MODELS, y_values = overall_speedups, legend_labels = devices,
               x_label = 'Neural Network Model', y_label = 'Speedup', bar_width = bar_width, x_min = 0.8,
               x_max = (x_ticks[-1] + 1), y_min = 0, y_max = y_max_val,
               folder = args.output[0] + '/Speedup/Overall', filename = 'exp_overall_speedup', allow_legend = True)
